from django.shortcuts import render
from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
from llm_sparsity.utils.text_generation import run
import json

def index(request):
    context = {'name': request.GET.get('submitted-name')}

    return render(request, 'index.html', context)

@csrf_exempt
def double_number(request):
    if request.method == 'POST':
        data = json.loads(request.body)
        num = int(data.get('number', 0))
        result = num * 2
        return JsonResponse({'result': result})

@csrf_exempt
def get_llmres(request):
    if request.method == 'POST':
        data = json.loads(request.body)
        if data['selectValue']=='opt':
            data['model_name']= r'E:\实训\llm_sparsity\llm_sparsity\model\opt-125m'
        elif data['selectValue']=='llama':
            data['model_name'] = 'huggyllama/llama-13b'
        elif data['selectValue']=='gpt':
            data['model_name'] = 'gpt'
        data['cache_dir']='/NVME1/cxl/hf_cache/'
        if data['mode']=='heavy':
            data['recentratio'] = "0.5"
            data['stride'] = '1'
        elif data['mode']=='recent':
            data['heavyratio'] = '0.0'
            data['stride'] = '1'
        elif data['mode'] == 'stride':
            data['heavyratio'] = '0.0'
            data['recentratio'] = '0.0'

        if data['precision']=='fp16':
            data['fp16']=True
        else:
            data['fp16']=False
        # print("data", data)
        # res_ori = {}
        if data['fp16']:
            res_ori, res_opt=run(["--model_name",data['model_name'],"--cache_dir",data['cache_dir'],"--mode",data['mode'],"--stride_step",data['stride'],"--heavy_ratio",data['heavyratio'],"--recent_ratio",data['recentratio'],"--length",data['length'],"--fp16"],data['inputText'])
        else:
            res_ori, res_opt = run(
                ["--model_name", data['model_name'], "--cache_dir", data['cache_dir'],"--mode",data['mode'],"--stride_step",data['stride'], "--heavy_ratio",
                 data['heavyratio'], "--recent_ratio", data['recentratio'], "--length",
                 data['length']],data['inputText'])
        # res_ori = res_opt
        return JsonResponse({'ori': res_ori,'opt': res_opt})
